from typing import List, Optional
from fastapi import FastAPI, HTTPException, Depends
from sqlmodel import SQLModel, Field, Relationship, Session, create_engine, select
from contextlib import asynccontextmanager
from fastapi.responses import JSONResponse
from datetime import datetime
# 数据库连接配置
DATABASE_URL = "sqlite:///database.db"

def current_time_str():
    return datetime.now().strftime("%Y-%m-%d-%H:%M:%S")

# 中间关联表
class TrainingTaskDatasetLink(SQLModel, table=True):
    training_task_id: int = Field(foreign_key="trainingtask.id", primary_key=True)
    dataset_id: int = Field(foreign_key="dataset.id", primary_key=True)


class DatasetPromptTemplateLink(SQLModel, table=True):
    dataset_id: int = Field(foreign_key="dataset.id", primary_key=True)
    prompt_template_id: int = Field(foreign_key="prompttemplate.id", primary_key=True)


# 主数据表
class DatasetBase(SQLModel):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(index=True)
    description: str | None = Field(default="", description="数据集描述")
    created_at: str = Field(default_factory=current_time_str)
    updated_at: str = Field(default_factory=current_time_str)
    type: str = Field(default="csv",description="数据集类型: 'csv' 或 'json'")
    original_path: str | None = Field(default="", description="原始CSV/json文件存储路径")
    train_file_path: Optional[str] = Field(default=None, description="训练集文件路径")
    test_file_path: Optional[str] = Field(default=None, description="测试集文件路径")

class Dataset(DatasetBase, table=True):
    training_tasks: List["TrainingTask"] = Relationship(
        back_populates="datasets", link_model=TrainingTaskDatasetLink
    )
    prompt_templates: List["PromptTemplate"] = Relationship(
        back_populates="datasets", link_model=DatasetPromptTemplateLink
    )
    prediction_tasks: List["PredictionTask"] = Relationship(back_populates="dataset")
    evaluation_tasks: List["EvaluationTask"] = Relationship(back_populates="dataset")


class DatasetCreate(DatasetBase):
    pass
class DatasetRead(DatasetBase):
    id: int 
class DatasetNoRels(DatasetBase):
    id: int

class DatasetWithRels(DatasetRead):
    training_tasks: List["TrainingTaskNoRels"] = []
    prompt_templates: List["PromptTemplateNoRels"] = []


class TrainingTaskBase(SQLModel):
    name: str
    status: str = Field(default="pending")  
    created_at: str = Field(default_factory=current_time_str)
    completed_at: str = Field(default_factory=current_time_str)


class TrainingTask(TrainingTaskBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    models: List["Model"] = Relationship(back_populates="training_task")
    datasets: List[Dataset] = Relationship(
        back_populates="training_tasks", link_model=TrainingTaskDatasetLink
    )
class TrainingTaskCreate(TrainingTaskBase):
    dataset_ids: List[int] = []  # 创建时可直接指定关联的数据集ID列表

class TrainingTaskRead(TrainingTaskBase):
    id: int

class TrainingTaskNoRels(TrainingTaskBase):
    id: int


class TrainingTaskWithRels(TrainingTaskBase):
    id: int | None = None
    datasets: List["DatasetNoRels"] = []
    models: List["ModelNoRels"] = []


class PromptTemplateBase(SQLModel):
    content: str


class PromptTemplate(PromptTemplateBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    datasets: List[Dataset] = Relationship(
        back_populates="prompt_templates", link_model=DatasetPromptTemplateLink
    )
    
class PromptTemplateRead(PromptTemplateBase):
    id: int
class PromptTemplateNoRels(PromptTemplateBase):
    id: int

class ModelBase(SQLModel):
    modelName: str
    modelType: str
    path: str
    description: str  | None = None
    training_task_id: int | None = Field(default=None, foreign_key="trainingtask.id")

class Model(ModelBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    training_task: TrainingTask | None = Relationship(back_populates="models")
    deployment_tasks: List["DeploymentTask"] = Relationship(back_populates="model")
    prediction_tasks: List["PredictionTask"] = Relationship(back_populates="model")
    evaluation_tasks: List["EvaluationTask"] = Relationship(back_populates="model")


class ModelCreate(ModelBase):
    pass
class ModelRead(ModelBase):
    id: int
class ModelNoRels(ModelBase):
    id: int

class ModelWithRels(ModelRead):
    training_task: "TrainingTaskNoRels" = None   
    deployment_tasks: List["DeploymentTaskNoRels"] = []
    prediction_tasks: List["PredictionTaskRead"] = []
    evaluation_tasks: List["EvaluationTaskRead"] = []

class DeploymentTaskBase(SQLModel):
    name: str | None = "deployment" + current_time_str()
    endpoint: str | None = None
    pgid: int | None = None
    created_at: str | None = Field(default_factory=current_time_str)
    completed_at: str | None = Field(default_factory=current_time_str)
    model_id: int | None = Field(default=None, foreign_key="model.id")


class DeploymentTask(DeploymentTaskBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    model: Model = Relationship(back_populates="deployment_tasks")

class DeploymentTaskCreate(DeploymentTaskBase):
    pass
class DeploymentTaskRead(DeploymentTaskBase):
    id: int

class DeploymentTaskWithRels(DeploymentTaskBase):
    id: int
    model: ModelNoRels | None = None
class DeploymentTaskNoRels(DeploymentTaskBase):
    id: int


class PredictionTaskBase(SQLModel):
    status: str
    model_id: int | None = Field(default=None, foreign_key="model.id")
    dataset_id: int | None = Field(default=None, foreign_key="dataset.id")


class PredictionTask(PredictionTaskBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    model: Model = Relationship(back_populates="prediction_tasks")
    dataset: Dataset = Relationship(back_populates="prediction_tasks")

class PredictionTaskRead(PredictionTaskBase):
    id: int
class PredictionTaskWithRels(PredictionTaskBase):
    id: int
    model: ModelNoRels | None = None
    dataset: DatasetNoRels | None = None


class EvaluationTaskBase(SQLModel):
    score: float
    model_id: int | None = Field(default=None, foreign_key="model.id")
    dataset_id: int | None = Field(default=None, foreign_key="dataset.id")


class EvaluationTask(EvaluationTaskBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    model: Model = Relationship(back_populates="evaluation_tasks")
    dataset: Dataset = Relationship(back_populates="evaluation_tasks")

class EvaluationTaskRead(EvaluationTaskBase):
    id: int

# 批量关联操作的请求模型
class BatchAssociateRequest(SQLModel):
    ids: List[int]


# 数据库引擎
engine = create_engine(DATABASE_URL, echo=True)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


@asynccontextmanager
async def lifespan(app: FastAPI):
    create_db_and_tables()
    yield


app = FastAPI(lifespan=lifespan)


def get_session():
    with Session(engine) as session:
        yield session


# Dataset 
@app.post("/datasets/", response_model=DatasetRead)
def create_dataset(dataset: DatasetCreate, session: Session = Depends(get_session)):
    db_dataset = Dataset(**dataset.model_dump())
    session.add(db_dataset)
    session.commit()
    session.refresh(db_dataset)
    return db_dataset


@app.get("/datasets/", response_model=List[DatasetWithRels])
def read_datasets(session: Session = Depends(get_session)):
    datasets = session.exec(select(Dataset)).all()
    return datasets


@app.get("/datasets/{dataset_id}", response_model=DatasetWithRels)
def read_dataset(dataset_id: int, session: Session = Depends(get_session)):
    dataset = session.get(Dataset, dataset_id)
    if not dataset:
        raise HTTPException(status_code=404, detail="Dataset not found")
    return dataset


# TrainingTask 
@app.post("/training_tasks/", response_model=TrainingTaskRead)
def create_training_task(
    task: TrainingTaskCreate, session: Session = Depends(get_session)
):
    # 提取数据集ID列表
    dataset_ids = task.dataset_ids
    task_data = task.model_dump(exclude={"dataset_ids"})

    # 创建训练任务
    db_task = TrainingTask(**task_data)
    session.add(db_task)
    session.commit()
    session.refresh(db_task)

    # 关联数据集
    if dataset_ids:
        for dataset_id in dataset_ids:
            dataset = session.get(Dataset, dataset_id)
            if dataset and dataset not in db_task.datasets:
                db_task.datasets.append(dataset)
        session.commit()
        session.refresh(db_task)

    return db_task

@app.get("/training_tasks/{task_id}", response_model=TrainingTaskWithRels)
def read_training_task(task_id: int, session: Session = Depends(get_session)):
    task = session.get(TrainingTask, task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Training task not found")
    return task

@app.get("/training_tasks/", response_model=List[TrainingTaskWithRels])
def read_training_tasks(session: Session = Depends(get_session)):
    tasks = session.exec(select(TrainingTask)).all()
    return tasks

# dataset关联操作
@app.post("/datasets/{dataset_id}/add_training_task/{task_id}")
def add_training_task_to_dataset(
    dataset_id: int, task_id: int, session: Session = Depends(get_session)
):
    dataset = session.get(Dataset, dataset_id)
    task = session.get(TrainingTask, task_id)

    if not dataset or not task:
        raise HTTPException(status_code=404, detail="Resource not found")

    if task not in dataset.training_tasks:
        dataset.training_tasks.append(task)
        session.commit()

    return {"message": "Training task added to dataset"}

# 批量关联数据集到训练任务
@app.post("/training_tasks/{task_id}/add_datasets")
def add_datasets_to_training_task(
    task_id: int, request: BatchAssociateRequest, session: Session = Depends(get_session)
):
    task = session.get(TrainingTask, task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Training task not found")

    added_count = 0
    for dataset_id in request.ids:
        dataset = session.get(Dataset, dataset_id)
        if dataset and dataset not in task.datasets:
            task.datasets.append(dataset)
            added_count += 1

    session.commit()
    return {"message": f"Added {added_count} datasets to training task"}

# 批量关联训练任务到数据集
@app.post("/datasets/{dataset_id}/add_training_tasks")
def add_training_tasks_to_dataset(
    dataset_id: int, request: BatchAssociateRequest, session: Session = Depends(get_session)
):
    dataset = session.get(Dataset, dataset_id)
    if not dataset:
        raise HTTPException(status_code=404, detail="Dataset not found")

    added_count = 0
    for task_id in request.ids:
        task = session.get(TrainingTask, task_id)
        if task and task not in dataset.training_tasks:
            dataset.training_tasks.append(task)
            added_count += 1

    session.commit()
    return {"message": f"Added {added_count} training tasks to dataset"}

@app.post("/datasets/{dataset_id}/add_prompt_template/{template_id}")
def add_prompt_template_to_dataset(
    dataset_id: int, template_id: int, session: Session = Depends(get_session)
):
    dataset = session.get(Dataset, dataset_id)
    template = session.get(PromptTemplate, template_id)

    if not dataset or not template:
        raise HTTPException(status_code=404, detail="Resource not found")

    if template not in dataset.prompt_templates:
        dataset.prompt_templates.append(template)
        session.commit()

    return {"message": "Prompt template added to dataset"}

# model
@app.post("/models/", response_model=ModelRead)
def create_model(model: ModelCreate, session: Session = Depends(get_session)):
    # 检查关联的训练任务是否存在
    task = session.get(TrainingTask, model.training_task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Training task not found")

    db_model = Model(**model.model_dump())
    session.add(db_model)
    session.commit()
    session.refresh(db_model)
    return db_model


@app.get("/models/", response_model=List[ModelWithRels])
def read_models(session: Session = Depends(get_session)):
    models = session.exec(select(Model)).all()
    return models

@app.get("/models/{model_id}", response_model=ModelWithRels)
def read_model(model_id: int, session: Session = Depends(get_session)):
    model = session.get(Model, model_id)
    if not model:
        return JSONResponse({"error": "Model not found"}, status_code=404)   
    return model

# deployment task
@app.post("/deployment_tasks/", response_model=DeploymentTaskRead)
def create_deployment_task(
    task: DeploymentTaskCreate, session: Session = Depends(get_session)
):
    db_task = DeploymentTask.model_validate(task)
    session.add(db_task)
    session.commit()
    session.refresh(db_task)
    return db_task

@app.get("/deployment_tasks/", response_model=List[DeploymentTaskWithRels])
def read_deployment_tasks(session: Session = Depends(get_session)):
    tasks = session.exec(select(DeploymentTask)).all()
    return tasks

@app.get("/deployment_tasks/{task_id}", response_model=DeploymentTaskWithRels)
def read_deployment_task(task_id: int, session: Session = Depends(get_session)):
    task = session.get(DeploymentTask, task_id)
    if not task:
        return  JSONResponse({"error": "Deployment task not found"}, status_code=404)
    return task
    

    
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=39999)